{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import sqlparse\n",
    "import pickle as pkl\n",
    "dataset_names = ['academic', 'atis', 'advising', 'geography', 'imdb', 'restaurants', 'scholar', 'yelp']\n",
    "\n",
    "# these datasets are small, so we use the full set. \n",
    "new_split_defined = {'restaurants', 'academic', 'imdb', 'yelp'} "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loading the original datasets from the paper:\n",
    "# Improving Text-to-SQL Evaluation Methodology\n",
    "\n",
    "# a dataset is a list of dictionaries\n",
    "# in the original dictionary, each datapoint might consist of several natural language sentences or SQL\n",
    "orig_datasets = []\n",
    "for dataset_name in dataset_names:\n",
    "    orig_dataset = json.load(open('text2sql-data/data/%s.json' % dataset_name))\n",
    "    for idx, d in enumerate(orig_dataset):\n",
    "        \n",
    "        d['orig_id'] = (dataset_name, idx)\n",
    "        \n",
    "        # fixing annotations here\n",
    "        \n",
    "        # change \"company_name\" to producer name, otherwise there is no variable to replace\n",
    "        if dataset_name == 'imdb' and idx == 27:\n",
    "            d['sql'][0] = 'SELECT MOVIEalias0.TITLE FROM COMPANY AS COMPANYalias0 , COPYRIGHT AS COPYRIGHTalias0 , MOVIE AS MOVIEalias0 WHERE COMPANYalias0.NAME = \"producer_name0\" AND COPYRIGHTalias0.CID = COMPANYalias0.ID AND MOVIEalias0.MID = COPYRIGHTalias0.MSID AND MOVIEalias0.RELEASE_YEAR > movie_release_year0 ;'\n",
    "    \n",
    "        # removing the extra space surrounding the variable actor_name0\n",
    "        if dataset_name == 'imdb' and idx == 78:\n",
    "            d['sql'][0] = 'SELECT MAX( DERIVED_TABLEalias0.DERIVED_FIELDalias0 ) FROM ( SELECT COUNT( DISTINCT ( MOVIEalias0.TITLE ) ) AS DERIVED_FIELDalias0 FROM ACTOR AS ACTORalias0 , CAST AS CASTalias0 , MOVIE AS MOVIEalias0 WHERE ACTORalias0.NAME = \"actor_name0\" AND CASTalias0.AID = ACTORalias0.AID AND MOVIEalias0.MID = CASTalias0.MSID GROUP BY MOVIEalias0.RELEASE_YEAR ) AS DERIVED_TABLEalias0 ;'\n",
    "    \n",
    "        # there was a scoping error; changed AUTHORalias1 to AUTHORalias0, PUBLICATIONalias1 to PUBLICATIONalias0\n",
    "        if dataset_name == 'academic' and idx == 182:\n",
    "            d['sql'][0] = 'SELECT DERIVED_FIELDalias0 FROM ( SELECT AUTHORalias0.NAME AS DERIVED_FIELDalias0 , COUNT( DISTINCT ( PUBLICATIONalias0.TITLE ) ) AS DERIVED_FIELDalias1 FROM AUTHOR AS AUTHORalias0 , CONFERENCE AS CONFERENCEalias0 , PUBLICATION AS PUBLICATIONalias0 , WRITES AS WRITESalias0 WHERE CONFERENCEalias0.NAME = \"conference_name0\" AND PUBLICATIONalias0.CID = CONFERENCEalias0.CID AND WRITESalias0.AID = AUTHORalias0.AID AND WRITESalias0.PID = PUBLICATIONalias0.PID GROUP BY AUTHORalias0.NAME ) AS DERIVED_TABLEalias0 , ( SELECT AUTHORalias1.NAME AS DERIVED_FIELDalias2 , COUNT( DISTINCT ( PUBLICATIONalias1.TITLE ) ) AS DERIVED_FIELDalias3 FROM AUTHOR AS AUTHORalias1 , CONFERENCE AS CONFERENCEalias1 , PUBLICATION AS PUBLICATIONalias1 , WRITES AS WRITESalias1 WHERE CONFERENCEalias1.NAME = \"conference_name1\" AND PUBLICATIONalias1.CID = CONFERENCEalias1.CID AND WRITESalias1.AID = AUTHORalias1.AID AND WRITESalias1.PID = PUBLICATIONalias1.PID GROUP BY AUTHORalias1.NAME ) AS DERIVED_TABLEalias1 WHERE DERIVED_TABLEalias0.DERIVED_FIELDalias1 > DERIVED_TABLEalias1.DERIVED_FIELDalias3 AND DERIVED_TABLEalias1.DERIVED_FIELDalias2 = DERIVED_TABLEalias0.DERIVED_FIELDalias0 ;'\n",
    "        \n",
    "        # wrong number of arguments to function COUNT(), change from \",\" to \"||\" for sqlite3 to recognize and execute\n",
    "        if dataset_name == 'advising' and idx == 107:\n",
    "            d['sql'][0] = 'SELECT COUNT( DISTINCT COURSEalias1.DEPARTMENT || COURSEalias0.NUMBER ) FROM COURSE AS COURSEalias0 , COURSE AS COURSEalias1 , COURSE_PREREQUISITE AS COURSE_PREREQUISITEalias0 , STUDENT_RECORD AS STUDENT_RECORDalias0 WHERE COURSEalias0.COURSE_ID = COURSE_PREREQUISITEalias0.PRE_COURSE_ID AND COURSEalias1.COURSE_ID = COURSE_PREREQUISITEalias0.COURSE_ID AND COURSEalias1.DEPARTMENT = \"department0\" AND COURSEalias1.NUMBER = number0 AND STUDENT_RECORDalias0.COURSE_ID = COURSEalias0.COURSE_ID AND STUDENT_RECORDalias0.STUDENT_ID = 1 ;'\n",
    "        \n",
    "        # there was not example given for level1 and hence replacing variable with values leads to errors\n",
    "        if dataset_name == 'advising' and idx == 132:\n",
    "            d['variables'][0]['example'] = '300'\n",
    "        \n",
    "        # cannot use count and order without group by; added grouping by actor_id\n",
    "        if dataset_name == 'imdb' and idx == 79:\n",
    "            d['sql'][0] = 'SELECT ACTORalias0.NAME FROM ACTOR AS ACTORalias0 , CAST AS CASTalias0 , MOVIE AS MOVIEalias0 WHERE CASTalias0.AID = ACTORalias0.AID AND MOVIEalias0.MID = CASTalias0.MSID GROUP BY ACTORalias0.AID ORDER BY COUNT( DISTINCT ( MOVIEalias0.TITLE ) ) DESC LIMIT 1 ;'\n",
    "    \n",
    "        # cannot use count and order without group by; added grouping by actor_id\n",
    "        if dataset_name == 'imdb' and idx == 80:\n",
    "            d['sql'][0] = 'SELECT ACTORalias0.NAME FROM ACTOR AS ACTORalias0 , CAST AS CASTalias0 , DIRECTED_BY AS DIRECTED_BYalias0 , DIRECTOR AS DIRECTORalias0 , MOVIE AS MOVIEalias0 WHERE CASTalias0.AID = ACTORalias0.AID AND DIRECTORalias0.DID = DIRECTED_BYalias0.DID AND MOVIEalias0.MID = CASTalias0.MSID AND MOVIEalias0.MID = DIRECTED_BYalias0.MSID GROUP BY ACTORalias0.AID ORDER BY COUNT( DISTINCT ( MOVIEalias0.TITLE ) ) DESC LIMIT 1 ;'\n",
    "        \n",
    "        # table has \"u\" in the neighborhood spelling.\n",
    "        n_before, n_after = 'NEIGHBORHOOD', 'NEIGHBOURHOOD'\n",
    "        if dataset_name == 'yelp':\n",
    "            d['sql'][0] = d['sql'][0].replace(n_before, n_after)\n",
    "        \n",
    "        if dataset_name == 'yelp' and idx == 42:\n",
    "            d['sql'][0] = 'SELECT NEIGHBOURHOODalias0.NEIGHBOURHOOD_NAME FROM BUSINESS AS BUSINESSalias0 , NEIGHBOURHOOD AS NEIGHBOURHOODalias0 , REVIEW AS REVIEWalias0 , USER AS USERalias0 WHERE NEIGHBOURHOODalias0.BUSINESS_ID = BUSINESSalias0.BUSINESS_ID AND REVIEWalias0.BUSINESS_ID = BUSINESSalias0.BUSINESS_ID AND USERalias0.NAME = \"user_name0\" AND USERalias0.USER_ID = REVIEWalias0.USER_ID ;'\n",
    "\n",
    "    orig_datasets.extend(orig_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are 3509 datapoints in the new testset\n"
     ]
    }
   ],
   "source": [
    "# we create the new testset here\n",
    "new_testset = []\n",
    "for d in orig_datasets:\n",
    "    orig_id = d['orig_id']\n",
    "    db_id, idx = orig_id\n",
    "    \n",
    "    # we only incorporate the test split if the dataset is large enough\n",
    "    # otherwise we incorporate the entire dataset\n",
    "    if d['query-split'] != 'test' and db_id not in new_split_defined:\n",
    "        continue\n",
    "    sql = d['sql'][0]\n",
    "    instance_variables = d['variables']\n",
    "    instance_name2examples = {d['name']: d['example'] for d in instance_variables}\n",
    "    \n",
    "    # we create a new datapoint for each natural language query\n",
    "    for sentence in d['sentences']:\n",
    "        new_datapoint = {\n",
    "            'text': sentence['text'],\n",
    "            'query': sql,\n",
    "            'variables': instance_variables,\n",
    "            'orig_id': orig_id,\n",
    "            'db_id': db_id,\n",
    "            'db_path': 'database/{db_id}/{db_id}.sqlite'.format(db_id=db_id)\n",
    "        }\n",
    "        new_testset.append(new_datapoint)\n",
    "print('There are %d datapoints in the new testset' % len(new_testset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "# this block implements a function that extract variable names from text and sql\n",
    "# later we use it to ensure that every variable is replaced\n",
    "\n",
    "variable_pattern = re.compile('^[a-z_]+[0-9]+$')\n",
    "\n",
    "def extract_variable_names(t):\n",
    "    tokens = t.replace('\"', '').replace('%', '').split(' ')\n",
    "    var_names = {v for v in tokens if variable_pattern.match(v) and 'alias' not in v}\n",
    "    return var_names\n",
    "\n",
    "test = False\n",
    "if test:\n",
    "    sql = 'SELECT BUSINESSalias0.NAME FROM BUSINESS AS BUSINESSalias0 , REVIEW AS REVIEWalias0 WHERE REVIEWalias0.BUSINESS_ID = BUSINESSalias0.BUSINESS_ID AND REVIEWalias0.MONTH = \"review_month0\" GROUP BY BUSINESSalias0.NAME ORDER BY COUNT( DISTINCT ( REVIEWalias0.TEXT ) ) DESC LIMIT 1 ;'\n",
    "    print(extract_variable_names(sql))\n",
    "    text = 'return me the homepage of journal_name0 .'\n",
    "    print(extract_variable_names(text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this block removes extra space surrounding variable names\n",
    "def remove_extra_space_around_variable(t):\n",
    "    var_names = extract_variable_names(t)\n",
    "    result = str(t)\n",
    "    for v in var_names:\n",
    "        result = result.replace('\" ' + v + ' \"', v)\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "set()\n"
     ]
    }
   ],
   "source": [
    "problematic = set()\n",
    "\n",
    "for datapoint in new_testset:\n",
    "    orig_id = datapoint['orig_id']\n",
    "    \n",
    "    # remove extra whitespace surrounding the text\n",
    "    datapoint['text'] = remove_extra_space_around_variable(datapoint['text'])\n",
    "    \n",
    "    # there should not be extra whitespace surrounding the sql variables\n",
    "    if datapoint['query'] != remove_extra_space_around_variable(datapoint['query']):\n",
    "        problematic.add(orig_id)\n",
    "\n",
    "    text_vars = extract_variable_names(datapoint['text'])\n",
    "    sql_vars = extract_variable_names(datapoint['query'])\n",
    "    \n",
    "    instance_variables = {d['name']: d for d in datapoint['variables']}\n",
    "    \n",
    "    # we ensure that all the variables in the sql query and the text can be replaced\n",
    "    # by some variable in the variable dictionary\n",
    "    if len(text_vars - instance_variables.keys()) != 0 or len(sql_vars - instance_variables.keys()):\n",
    "        problematic.add(orig_id)\n",
    "        \n",
    "    # replace the variables with the examples in the variable dictionary\n",
    "    for text_var in text_vars:\n",
    "        datapoint['text'] = datapoint['text'].replace(text_var, instance_variables[text_var]['example'])\n",
    "    \n",
    "    for sql_var in sql_vars:\n",
    "        datapoint['query'] = datapoint['query'].replace(sql_var, instance_variables[sql_var]['example'])\n",
    "\n",
    "# we can trace back which datapoints do not satisfy the assumption,\n",
    "# then go back and fix it manually\n",
    "print(problematic)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'db_id': 'academic',\n",
      "  'db_path': 'database/academic/academic.sqlite',\n",
      "  'orig_id': ('academic', 0),\n",
      "  'query': 'SELECT JOURNALalias0.HOMEPAGE FROM JOURNAL AS JOURNALalias0 WHERE '\n",
      "           'JOURNALalias0.NAME = \"PVLDB\" ;',\n",
      "  'text': 'return me the homepage of PVLDB .',\n",
      "  'variables': [{'example': 'PVLDB',\n",
      "                 'location': 'both',\n",
      "                 'name': 'journal_name0',\n",
      "                 'type': 'journal_name'}]},\n",
      " {'db_id': 'academic',\n",
      "  'db_path': 'database/academic/academic.sqlite',\n",
      "  'orig_id': ('academic', 1),\n",
      "  'query': 'SELECT AUTHORalias0.HOMEPAGE FROM AUTHOR AS AUTHORalias0 WHERE '\n",
      "           'AUTHORalias0.NAME = \"H. V. Jagadish\" ;',\n",
      "  'text': 'return me the homepage of H. V. Jagadish .',\n",
      "  'variables': [{'example': 'H. V. Jagadish',\n",
      "                 'location': 'both',\n",
      "                 'name': 'author_name0',\n",
      "                 'type': 'author_name'}]},\n",
      " {'db_id': 'academic',\n",
      "  'db_path': 'database/academic/academic.sqlite',\n",
      "  'orig_id': ('academic', 2),\n",
      "  'query': 'SELECT PUBLICATIONalias0.ABSTRACT FROM PUBLICATION AS '\n",
      "           'PUBLICATIONalias0 WHERE PUBLICATIONalias0.TITLE = \"Making database '\n",
      "           'systems usable\" ;',\n",
      "  'text': 'return me the abstract of Making database systems usable .',\n",
      "  'variables': [{'example': 'Making database systems usable',\n",
      "                 'location': 'both',\n",
      "                 'name': 'publication_title0',\n",
      "                 'type': 'publication_title'}]},\n",
      " {'db_id': 'academic',\n",
      "  'db_path': 'database/academic/academic.sqlite',\n",
      "  'orig_id': ('academic', 3),\n",
      "  'query': 'SELECT PUBLICATIONalias0.YEAR FROM PUBLICATION AS '\n",
      "           'PUBLICATIONalias0 WHERE PUBLICATIONalias0.TITLE = \"Making database '\n",
      "           'systems usable\" ;',\n",
      "  'text': 'return me the year of Making database systems usable',\n",
      "  'variables': [{'example': 'Making database systems usable',\n",
      "                 'location': 'both',\n",
      "                 'name': 'publication_title0',\n",
      "                 'type': 'publication_title'}]},\n",
      " {'db_id': 'academic',\n",
      "  'db_path': 'database/academic/academic.sqlite',\n",
      "  'orig_id': ('academic', 3),\n",
      "  'query': 'SELECT PUBLICATIONalias0.YEAR FROM PUBLICATION AS '\n",
      "           'PUBLICATIONalias0 WHERE PUBLICATIONalias0.TITLE = \"Making database '\n",
      "           'systems usable\" ;',\n",
      "  'text': 'return me the year of Making database systems usable .',\n",
      "  'variables': [{'example': 'Making database systems usable',\n",
      "                 'location': 'both',\n",
      "                 'name': 'publication_title0',\n",
      "                 'type': 'publication_title'}]}]\n"
     ]
    }
   ],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "pprint(new_testset[:5])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
